import click
import rasterio
from rasterio.windows import Window
from concurrent.futures import ThreadPoolExecutor, as_completed
import os
import sys
import logging
import mercantile
import math
from datetime import datetime

# 设置环境变量以抑制 PROJ 警告
os.environ['PROJ_LIB'] = os.path.join(os.path.dirname(rasterio.__file__), 'proj_data')
os.environ['GDAL_DATA'] = os.path.join(os.path.dirname(rasterio.__file__), 'gdal_data')

# 配置日志系统
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    handlers=[logging.FileHandler("tile_worker.log"), logging.StreamHandler()]
)
logger = logging.getLogger(__name__)

def gdal_translate(input_path, output_path, threads):
    """执行GDAL转换命令"""
    cmd = (
        f'gdal_translate -of MBTiles "{input_path}" "{output_path}" '
        f'-co TILE_FORMAT=PNG -co QUALITY=90 '
        f'-co ZOOM_LEVEL_STRATEGY=auto '
        f'--config GDAL_NUM_THREADS {threads} '
        f'--config GDAL_FORCE_SQLITE_JOURNAL OFF '
        f'--config OGR_SQLITE_SYNCHRONOUS OFF'
    )
    logger.info(f"执行GDAL命令: {cmd}")
    exit_code = os.system(cmd)
    if exit_code != 0:
        raise RuntimeError(f"GDAL执行失败，错误码: {exit_code}")

def generate_tiles(bounds, min_zoom, max_zoom):
    """使用mercantile生成覆盖范围的瓦片列表"""
    for zoom in range(min_zoom, max_zoom + 1):
        for tile in mercantile.tiles(
            west=bounds['west'],
            south=bounds['south'],
            east=bounds['east'],
            north=bounds['north'],
            zooms=zoom
        ):
            yield tile

def process_tile(src, output_path, tile):
    """处理单个瓦片并写入MBTiles"""
    try:
        # 将TMS坐标转换为GDAL使用的XYZ格式
        xyz_tile = mercantile.Tile(tile.x, (1 << tile.z) - 1 - tile.y, tile.z)
        
        # 计算像素窗口
        transform = src.transform
        height = src.height
        
        # 计算瓦片边界坐标
        ul = mercantile.ul(tile)
        lr = mercantile.ul(mercantile.Tile(tile.x + 1, tile.y + 1, tile.z))
        
        # 转换为像素坐标
        px_left, px_top = ~transform * (ul[0], ul[1])
        px_right, px_bottom = ~transform * (lr[0], lr[1])
        
        # 确保坐标在有效范围内
        px_left = max(0, min(px_left, src.width))
        px_right = max(0, min(px_right, src.width))
        px_top = max(0, min(px_top, height))
        px_bottom = max(0, min(px_bottom, height))
        
        # 创建窗口
        window = Window.from_slices(
            rows=(int(px_top), int(px_bottom)),
            cols=(int(px_left), int(px_right))
        )
        
        # 读取数据并写入MBTiles
        with rasterio.open(output_path, "r+") as dst:
            data = src.read(window=window)
            if data.size > 0:  # 确保有数据要写入
                dst.write(data, window=Window.from_slices(
                    (xyz_tile.y * 256, (xyz_tile.y + 1) * 256),
                    (xyz_tile.x * 256, (xyz_tile.x + 1) * 256))
                )
            
        return True
    except Exception as e:
        logger.error(f"瓦片 {tile.z}-{tile.x}-{tile.y} 处理失败: {str(e)}")
        raise

@click.command()
@click.option("--input", required=True, help="输入TIFF路径")
@click.option("--output", required=True, help="输出MBTiles路径")
@click.option("--threads", default=4, help="并行线程数")
@click.option("--min-zoom", default=0, help="最小缩放级别")
@click.option("--max-zoom", default=14, help="最大缩放级别")
def main(input, output, threads, min_zoom, max_zoom):
    """完整的MBTiles生成脚本"""
    start_time = datetime.now()
    
    try:
        # ===== 阶段1: 生成基础MBTiles =====
        logger.info("正在生成MBTiles基础结构...")
        gdal_translate(input, output, threads)
        
        # ===== 阶段2: 并行处理瓦片 =====
        with rasterio.open(input) as src:
            bounds = {
                'west': src.bounds.left,
                'south': src.bounds.bottom,
                'east': src.bounds.right,
                'north': src.bounds.top
            }
            
            # 生成瓦片列表
            logger.info("正在计算瓦片范围...")
            tiles = list(generate_tiles(bounds, min_zoom, max_zoom))
            logger.info(f"共需处理 {len(tiles)} 个瓦片")
            
            # 启动线程池
            with ThreadPoolExecutor(max_workers=threads) as executor:
                futures = []
                logger.info(f"启动线程池 ({threads} workers)")
                
                # 提交任务
                for idx, tile in enumerate(tiles):
                    future = executor.submit(process_tile, src, output, tile)
                    futures.append(future)
                    if idx % 1000 == 0:
                        logger.debug(f"已提交 {idx} 个任务")
                
                # 进度监控
                completed = 0
                for future in as_completed(futures):
                    try:
                        future.result()
                        completed += 1
                        if completed % 100 == 0:
                            progress = completed / len(futures) * 100
                            logger.info(
                                f"进度: {progress:.1f}% "
                                f"({completed}/{len(futures)})"
                            )
                    except Exception as e:
                        logger.error(f"任务失败: {str(e)}")
        
        # ===== 阶段3: 生成金字塔 =====
        logger.info("正在生成金字塔...")
        os.system(f'gdaladdo -r average "{output}" 2 4 8 16 32 64 128 256')
        
    except Exception as e:
        logger.error(f"处理失败: {str(e)}", exc_info=True)
        sys.exit(1)
        
    finally:
        logger.info(f"总耗时: {datetime.now() - start_time}")

if __name__ == "__main__":
    main()